/*
 * Copyright 2010-2024 JetBrains s.r.o. and Kotlin Programming Language contributors.
 * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
 */
package com.google.j2cl.transpiler.frontend.kotlin.lower

import org.jetbrains.kotlin.DeprecatedForRemovalCompilerApi
import org.jetbrains.kotlin.backend.common.BodyLoweringPass
import org.jetbrains.kotlin.backend.common.LoweringContext
import org.jetbrains.kotlin.backend.common.ir.isInlineFunWithReifiedParameter
import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
import org.jetbrains.kotlin.backend.common.lower.irBlock
import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
import org.jetbrains.kotlin.ir.builders.declarations.addValueParameter
import org.jetbrains.kotlin.ir.builders.declarations.buildFun
import org.jetbrains.kotlin.ir.builders.irCall
import org.jetbrains.kotlin.ir.builders.irGet
import org.jetbrains.kotlin.ir.builders.irReturn
import org.jetbrains.kotlin.ir.declarations.*
import org.jetbrains.kotlin.ir.expressions.IrBody
import org.jetbrains.kotlin.ir.expressions.IrExpression
import org.jetbrains.kotlin.ir.expressions.IrFunctionReference
import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
import org.jetbrains.kotlin.ir.expressions.impl.IrFunctionReferenceImpl
import org.jetbrains.kotlin.ir.expressions.impl.fromSymbolOwner
import org.jetbrains.kotlin.ir.symbols.IrTypeParameterSymbol
import org.jetbrains.kotlin.ir.types.IrTypeArgument
import org.jetbrains.kotlin.ir.types.IrTypeSubstitutor
import org.jetbrains.kotlin.ir.util.SYNTHETIC_OFFSET
import org.jetbrains.kotlin.ir.util.render
import org.jetbrains.kotlin.ir.util.setDeclarationsParent
import org.jetbrains.kotlin.ir.util.typeSubstitutionMap
import org.jetbrains.kotlin.ir.visitors.IrTransformer
import org.jetbrains.kotlin.name.Name
import org.jetbrains.kotlin.utils.addToStdlib.runIf

/**
 * Replaces callable reference to an inline function with reified parameter with a callable
 * reference to a new non-inline function with substituted types.
 *
 * Copied from
 * org/jetbrains/kotlin/backend/common/lower/WrapInlineDeclarationsWithReifiedTypeParametersLowering.kt
 *
 * TODO(b/449125803): Remove this pass in favor of UpgradeCallableReferences.
 */
@SuppressWarnings("CheckReturnValue")
@OptIn(DeprecatedForRemovalCompilerApi::class)
internal class WrapInlineDeclarationsWithReifiedTypeParametersLowering(
  val context: LoweringContext
) : BodyLoweringPass {
  private val irFactory
    get() = context.irFactory

  override fun lower(irBody: IrBody, container: IrDeclaration) {
    irBody.transformChildren(
      object : IrTransformer<IrDeclarationParent?>() {
        override fun visitDeclaration(declaration: IrDeclarationBase, data: IrDeclarationParent?) =
          super.visitDeclaration(declaration, declaration as? IrDeclarationParent ?: data)

        override fun visitFunctionReference(
          expression: IrFunctionReference,
          data: IrDeclarationParent?,
        ): IrExpression {
          expression.transformChildren(this, data)

          val owner = expression.symbol.owner as? IrSimpleFunction ?: return expression

          if (!owner.isInlineFunWithReifiedParameter()) {
            return expression
          }
          @Suppress("UNCHECKED_CAST")
          val typeSubstitutor =
            IrTypeSubstitutor(
              expression.typeSubstitutionMap as Map<IrTypeParameterSymbol, IrTypeArgument>
            )

          val function =
            irFactory
              .buildFun {
                name = Name.identifier("${owner.name}${"$"}wrap")
                returnType = typeSubstitutor.substitute(owner.returnType)
                visibility = DescriptorVisibilities.LOCAL
                origin = IrDeclarationOrigin.ADAPTER_FOR_CALLABLE_REFERENCE
                startOffset = SYNTHETIC_OFFSET
                endOffset = SYNTHETIC_OFFSET
              }
              .apply {
                parent =
                  data
                    ?: error(
                      "Unable to get a proper parent while lower ${expression.render()} at ${container.render()}"
                    )
                val irBuilder = context.createIrBuilder(symbol, SYNTHETIC_OFFSET, SYNTHETIC_OFFSET)
                val forwardExtensionReceiverAsParam =
                  owner.extensionReceiverParameter?.let { extensionReceiver ->
                    runIf(expression.extensionReceiver == null) {
                      addValueParameter(
                        extensionReceiver.name,
                        typeSubstitutor.substitute(extensionReceiver.type),
                      )
                      true
                    }
                  } ?: false
                owner.valueParameters.forEach { valueParameter ->
                  addValueParameter(
                    valueParameter.name,
                    typeSubstitutor.substitute(valueParameter.type),
                  )
                }
                body =
                  irFactory.createBlockBody(expression.startOffset, expression.endOffset) {
                    statements.add(
                      irBuilder.irReturn(
                        irBuilder.irCall(owner.symbol).also { call ->
                          expression.extensionReceiver?.setDeclarationsParent(this@apply)
                          expression.dispatchReceiver?.setDeclarationsParent(this@apply)
                          val (extensionReceiver, forwardedParams) =
                            if (forwardExtensionReceiverAsParam) {
                              irBuilder.irGet(valueParameters.first()) to
                                valueParameters.subList(1, valueParameters.size)
                            } else {
                              expression.extensionReceiver to valueParameters
                            }
                          call.extensionReceiver = extensionReceiver
                          call.dispatchReceiver = expression.dispatchReceiver

                          forwardedParams.forEachIndexed { index, valueParameter ->
                            call.putValueArgument(index, irBuilder.irGet(valueParameter))
                          }
                          for (i in expression.typeArguments.indices) {
                            call.typeArguments[i] = expression.typeArguments[i]
                          }
                        }
                      )
                    )
                  }
              }
          return context.createIrBuilder(container.symbol).irBlock(
            expression,
            origin = IrStatementOrigin.ADAPTED_FUNCTION_REFERENCE,
          ) {
            +function
            +IrFunctionReferenceImpl.fromSymbolOwner(
              expression.startOffset,
              expression.endOffset,
              expression.type,
              function.symbol,
              function.typeParameters.size,
              expression.reflectionTarget,
            )
          }
        }
      },
      container as? IrDeclarationParent,
    )
  }
}
